package server

import (
	"context"
	"github.com/sirupsen/logrus"
	"google.golang.org/grpc"
	"log"
	"net"
	"net/http"
	"simPSapp/pb"
	"strconv"
	"sync"
	"sync/atomic"
	"time"
)

type Server struct {
	grpcServer *grpc.Server
	svc        *Service

	health http.Server

	port string
}

type Service struct {
	pb.UnimplementedSendDataServer

	allDataSize   int
	allDataSizeMB int
	workerNum     int
	curDataSize   atomic.Int32
	dataBucket    []int // 存储每个worker已经上传的数据大小

	startTime *time.Time // 任务开始时间。lazy初始化当第一个请求到来时进行设置，因此设置为指针
	endTime   time.Time  // 所有数据收集完成后

	lockForBatchBarrier sync.RWMutex
	batchBarrier        map[int32]*sync.WaitGroup // 记录有多少个worker完成了当前批次的上传

	lockForStartSync sync.Mutex
	startSync        *sync.WaitGroup
}

type LivenessProbeService struct {
	pb.UnimplementedHealthServer
}

func (s *LivenessProbeService) Check(context.Context, *pb.HealthCheckRequest) (*pb.HealthCheckResponse, error) {
	return &pb.HealthCheckResponse{Status: pb.HealthCheckResponse_SERVING}, nil
}

func NewServer(allDataSize, workerNum int, port string) *Server {
	var serverOptions = []grpc.ServerOption{
		grpc.MaxRecvMsgSize(1024 * 1024 * 1024),
	}
	s := grpc.NewServer(serverOptions...)
	svc := &Service{
		allDataSize:   allDataSize,
		workerNum:     workerNum,
		dataBucket:    make([]int, workerNum),
		curDataSize:   atomic.Int32{},
		startTime:     nil,
		allDataSizeMB: allDataSize * 1024 * 1024,
		batchBarrier:  make(map[int32]*sync.WaitGroup),
	}
	pb.RegisterSendDataServer(s, svc)

	probeSvc := &LivenessProbeService{}
	pb.RegisterHealthServer(s, probeSvc)

	_, err := strconv.Atoi(port)
	if err != nil {
		logrus.Fatalf("端口号必须是数字: %v", err)
	}
	return &Server{grpcServer: s, svc: svc, port: port}

}

func (s *Service) SendNByte(ctx context.Context, r *pb.RequestBytes) (*pb.Response, error) {
	if s.startTime == nil {
		now := time.Now()
		s.startTime = &now
	}

	l := len(r.Payload)
	s.dataBucket[r.From] += l
	s.curDataSize.Add(int32(l))

	// check是否收集了所有数据
	if s.curDataSize.Load() >= int32(s.allDataSizeMB) {
		s.endTime = time.Now()
		logrus.Infof("任务完成，总耗时: %v", s.endTime.Sub(*s.startTime))
	}

	batchId := r.BatchId

	// 初始化wg
	s.lockForBatchBarrier.Lock()
	if s.batchBarrier[batchId] == nil {
		wg := &sync.WaitGroup{}
		wg.Add(s.workerNum)
		s.batchBarrier[batchId] = wg
	}
	s.lockForBatchBarrier.Unlock()

	s.batchBarrier[batchId].Done()

	// 等待所有worker上传完毕后，这个Wait会放行
	s.batchBarrier[batchId].Wait()

	return &pb.Response{Msg: "success"}, nil
}

func (s *Service) WorkerSync(ctx context.Context, r *pb.HelloRequest) (*pb.Response, error) {
	s.lockForStartSync.Lock()
	if s.startSync == nil {
		s.startSync = &sync.WaitGroup{}
		s.startSync.Add(s.workerNum)
	}
	s.lockForStartSync.Unlock()

	s.startSync.Done()

	s.startSync.Wait()

	return &pb.Response{Msg: "continue"}, nil
}

func (server *Server) Run() {
	port := server.port
	netListener, err := net.Listen("tcp", ":"+port)
	if err != nil {
		logrus.Fatalf("无法监听 端口(:%s) : %v", port, err)
	}

	logrus.Infof("grpc server will listen on port %s", port)

	_ = server.grpcServer.Serve(netListener)
}

func (server *Server) Stop() {
	server.grpcServer.Stop()
}

func WaitForGrpc(sv *Server) {
	for {
		conn, err := net.Dial("tcp", ":"+sv.port)
		if err == nil {
			conn.Close()
			log.Println("gRPC server is ready")
			return
		}
		log.Println("Waiting for gRPC server to start...")
		time.Sleep(1 * time.Second)
	}
}

//func (server *Server) unlock() {
//	server.svc.lockForStartSync.Unlock()
//	server.svc.lockForBatchBarrier.Unlock()
//}
